import matplotlib.pyplot as plt
from fileRW import ReadFile
from pre_handle import PreHandle
from backTest import BackTest
from visualization import Visualization
from resultAnalysis import ResultAnalysis
import numpy as np
from train_model import TrainModel


if __name__ == '__main__':
    initial_money = 1000000
    start_date = '2013-01-01'
    end_date = '2019-12-31'
    file_path = "./data.pkl"
    bt = BackTest(file_path="./data.pkl", start_date=start_date, end_date=end_date)
    res = bt.back_trader(strategy_name='涨幅策略')
    res_avg = bt.back_trader(strategy_name='双均线策略')

    days = list(res.index)
    return_list = list(res.values)
    return_list_avg = list(res_avg.values)

    for i in range(len(return_list)):
        return_list[i] = initial_money * (1 + return_list[i])
    for i in range(len(return_list_avg)):
        return_list_avg[i] = initial_money * (1 + return_list_avg[i])

    vs = Visualization(initial=initial_money)
    vs.plot_return(return_list, isSSEC=True, isCSI300=True, start_date=start_date, end_date=end_date, strategy='涨幅策略')
    vs.plot_return(return_list_avg, isSSEC=True, isCSI300=True, start_date=start_date, end_date=end_date, strategy='双均线策略')

    # 计算各项指标
    rs = ResultAnalysis(initial=initial_money, start_date=start_date, end_date=end_date)
    days, return_list_SSEC = rs.cal_return_list_SSEC()

    # print(type(return_list))
    # print(type())
    # 计算总收益率
    gain_ratio = rs.cal_gain_ratio(return_list)
    gain_ratio_avg = rs.cal_gain_ratio(return_list_avg)
    gain_ratio_SSEC = rs.cal_gain_ratio(return_list_SSEC)

    # 年化收益率
    gain_ratio_a_year = rs.annualized_rate_of_return(days, return_list)
    gain_ratio_a_year_avg = rs.annualized_rate_of_return(days, return_list_avg)
    gain_ratio_a_year_SSEC = rs.annualized_rate_of_return(days, return_list_SSEC)

    # 最大回撤率
    drawdown_rate, drawdown_max, drawdown_tian, _, _ = rs.cal_MaxDrawdown(return_list)
    drawdown_rate_avg, drawdown_max_avg, drawdown_tian_avg, _, _ = rs.cal_MaxDrawdown(return_list_avg)
    drawdown_rate_SSEC, drawdown_max_SSEC, drawdown_tian_SSEC, _, _ = rs.cal_MaxDrawdown(return_list_SSEC)

    # 夏普比率
    sharp_ratio = rs.cal_sharp_ratio(days, return_list)
    sharp_ratio_avg = rs.cal_sharp_ratio(days, return_list_avg)
    sharp_ratio_SSEC = rs.cal_sharp_ratio(days, return_list_SSEC)

    # 击败基准比率（日胜率）
    # beat_benchmark_ratio = rs.beat_benchmark_ratio(return_list)
    # beat_benchmark_ratio_SSEC = rs.beat_benchmark_ratio(return_list_SSEC)

    # 绘制分组条形图
    labels = ['总收益率', '年化收益率', '夏普比率', '最大回撤率']
    strategy1_score = np.array([gain_ratio, gain_ratio_a_year, sharp_ratio, drawdown_rate])
    strategy1_score = list(np.round(strategy1_score, 2))
    strategy2_score = np.array([gain_ratio_avg, gain_ratio_a_year_avg, sharp_ratio_avg, drawdown_rate_avg])
    strategy2_score = list(np.round(strategy2_score, 2))
    SSEC_score = np.array([gain_ratio_SSEC, gain_ratio_a_year_SSEC, sharp_ratio_SSEC, drawdown_rate_SSEC])
    SSEC_score = list(np.round(SSEC_score, 2))

    Visualization.group_bar_chart(labels, strategy1_score, SSEC_score, strategy='涨幅策略')

    Visualization.group_3bar_chart(labels, strategy1_score, strategy2_score, SSEC_score,
                                   strategy1='涨幅策略', strategy2='双均线策略')

